"""
2.	使用Python和OpenCV，借助背景减除算法进行运动车辆前景检测（36分）
"""
import cv2 as cv
import numpy as np
import os
import sys

HISTORY = 500
FPS = 25
INTERVAL = 1000 // FPS
ALPHA = 1e-3
BASE_DIR, FILE_NAME = os.path.split(__file__)
path = 'data/video/traffic_1min.mp4'
VIDEO_PATH = os.path.join(BASE_DIR, path)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME)
filename = os.path.join(SAVE_DIR, 'output.avi')

# ①	读取视频文件traffic_1min.mp4，提取所有背景帧
print('Training ...')
bg = cv.createBackgroundSubtractorMOG2(HISTORY, detectShadows=False)
video = cv.VideoCapture(VIDEO_PATH)
cnt = 0
while True:
    ret, img = video.read()
    if not ret:
        break
    cnt += 1
    if cnt > HISTORY:
        break
    bg.apply(img, None, ALPHA)
video.release()
print('Trained.')

# ②	使用形态学处理，比如背景消除、做帧差、膨胀腐蚀等
# ③	将处理的数据帧保存在自建文件夹里
# ④	文件夹里的所有图像帧，保存为视频文件并可播放

video = cv.VideoCapture(VIDEO_PATH)
fourcc = cv.VideoWriter_fourcc(*list('XVID'))
H, W = img.shape[:2]
os.makedirs(os.path.split(filename)[0], exist_ok=True)
output = cv.VideoWriter(filename, fourcc, FPS, (W, H))
print(f'Wrinting video to {filename}')
cnt = 0
kernel = np.ones((3, 3), dtype=np.uint8)
while True:
    ret, img = video.read()
    if not ret:
        print('Video over.')
        break
    cnt += 1

    mask = bg.apply(img, None, 0.)
    mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel)
    mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel)
    mask = cv.dilate(mask, kernel, iterations=2)

    img = cv.bitwise_and(img, img, mask=mask)
    cv.imshow('Traffic', img)
    output.write(img)

    k = cv.waitKey(INTERVAL) & 0xFF
    if 27 == k:
        print('Interrupted by user.')
        break
print(f'Processed {cnt} frames.')
video.release()
output.release()
print('Video written.')
print('Over')
